"""
暗区检测器实现类

基于分割模型实现暗区检测功能，支持单张图片和批量检测。
"""

import os
import sys
import numpy as np
from PIL import Image
from typing import Dict, Any, Optional, Union, List
from pathlib import Path

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference


class DarkDetectorImpl:
    """
    暗区检测器实现类
    
    基于分割模型实现暗区检测功能，可以检测图像中的暗区区域。
    """
    
    def __init__(self, 
                 model_path: str, 
                 device: str = 'cpu',
                 dark_class_id: int = 1,
                 confidence_threshold: float = 0.5,
                 area_threshold: int = 400000):
        """
        初始化暗区检测器
        
        Args:
            model_path (str): 分割模型文件路径
            device (str): 推理设备，'cpu' 或 'cuda'
            dark_class_id (int): 暗区类别ID，默认1
            confidence_threshold (float): 置信度阈值，默认0.5
            area_threshold (int): 面积阈值，默认400000像素点，大于此值认为有黑边
        """
        self.model_path = model_path
        self.device = device
        self.dark_class_id = dark_class_id
        self.confidence_threshold = confidence_threshold
        self.area_threshold = area_threshold
        
        # 创建推理器实例
        try:
            self.infer = create_inference(model_path, device=device)
            print(f"暗区检测器初始化成功，使用设备: {device}")
        except Exception as e:
            raise RuntimeError(f"初始化暗区检测器失败: {str(e)}")
    
    def detect_dark_single(self, image: Union[str, Image.Image, np.ndarray], 
                          output_path: Optional[str] = None) -> Dict[str, Any]:
        """
        检测单张图片中的暗区区域
        
        Args:
            image: 图像数据，可以是图像路径(str)、PIL Image对象或numpy数组
            output_path (str, optional): 输出图片路径，如果提供则保存带标注的图片
            
        Returns:
            Dict[str, Any]: 检测结果，包含暗区区域信息
        """
        try:
            # 处理不同类型的输入
            if isinstance(image, str):
                # 如果是路径，直接使用推理器
                result = self.infer.infer_single(image)
            elif isinstance(image, Image.Image):
                # 如果是PIL Image，直接使用推理器
                result = self.infer.infer_single(image)
            elif isinstance(image, np.ndarray):
                # 如果是numpy数组，直接传递给推理器（base_inference已处理格式转换）
                result = self.infer.infer_single(image)
            else:
                raise ValueError(f"不支持的图像类型: {type(image)}")
            
            # 提取暗区区域并进行过滤
            dark_regions, has_dark_edge = self._extract_dark_regions(result['segmentation_mask'])
            
            # 构建返回结果，包含完整的模型推理结果
            detection_result = {
                'dark_regions': dark_regions,
                'total_dark_regions': len(dark_regions),
                'has_dark_edge': has_dark_edge,
                'segmentation_mask': result['segmentation_mask'],
                'num_classes': result['num_classes'],
                'image_shape': result['segmentation_mask'].shape,
                # 添加完整的模型推理结果
                'model_result': result
            }
            
            # 如果提供了输出路径，保存结果图片
            if output_path:
                self._save_result_image(image, detection_result, output_path)
            
            return detection_result
            
        except Exception as e:
            raise RuntimeError(f"暗区检测失败: {str(e)}")
    
    def detect_dark_batch(self, input_dir: str) -> Dict[str, Any]:
        """
        批量检测目录中的图片暗区
        
        Args:
            input_dir (str): 输入图片目录
            
        Returns:
            Dict[str, Any]: 批量检测结果
        """
        try:
            # 使用推理器的批量推理功能
            results = self.infer.infer_batch(input_dir)
            
            # 处理每张图片的结果
            processed_results = {}
            for img_name, result in results.items():
                if 'error' in result:
                    processed_results[img_name] = {'error': result['error']}
                else:
                    # 提取暗区区域并进行过滤
                    dark_regions, has_dark_edge = self._extract_dark_regions(result['segmentation_mask'])
                    
                    processed_results[img_name] = {
                        'dark_regions': dark_regions,
                        'total_dark_regions': len(dark_regions),
                        'has_dark_edge': has_dark_edge,
                        'segmentation_mask': result['segmentation_mask'],
                        'num_classes': result['num_classes'],
                        'image_shape': result['segmentation_mask'].shape,
                        # 添加完整的模型推理结果
                        'model_result': result
                    }
            
            return processed_results
            
        except Exception as e:
            raise RuntimeError(f"批量暗区检测失败: {str(e)}")
    
    def _extract_dark_regions(self, segmentation_mask: np.ndarray) -> tuple[List[Dict[str, Any]], bool]:
        """
        从分割掩码中提取暗区区域，使用OpenCV findContours获取多边形轮廓
        只处理标签为1的黑边区域，忽略其他类别
        
        Args:
            segmentation_mask (np.ndarray): 分割掩码
            
        Returns:
            tuple[List[Dict[str, Any]], bool]: (暗区区域列表, 是否有黑边)
        """
        # 根据分割掩码尺寸动态调整面积阈值
        mask_height, mask_width = segmentation_mask.shape[:2]
        # 假设原始图像是2K-3K尺寸，计算缩放比例
        reference_size = 2500  # 参考尺寸（2K-3K的平均值）
        # 采用图像面积与参考面积的平方根比值作为缩放因子，避免长宽比过大导致不准确
        scale_factor = np.sqrt((mask_width * mask_height) / (reference_size * reference_size))
        # 计算调整后的面积阈值（不修改原始值）
        adjusted_area_threshold = int(self.area_threshold * (scale_factor ** 2))
        
        dark_regions = []
        has_dark_edge = False
        
        # 只找到标签为1的黑边像素（self.dark_class_id应该为1）
        dark_pixels = (segmentation_mask == self.dark_class_id)
        
        if not np.any(dark_pixels):
            return dark_regions, has_dark_edge
        
        # 转换为uint8格式用于OpenCV
        binary_mask = (dark_pixels * 255).astype(np.uint8)
        
        # 形态学操作来平滑边界
        import cv2
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)
        
        # 使用OpenCV findContours找到轮廓
        contours, hierarchy = cv2.findContours(
            binary_mask, 
            cv2.RETR_EXTERNAL,  # 只检测外部轮廓
            cv2.CHAIN_APPROX_NONE  # 保留所有轮廓点，不压缩
        )
        
        if not contours:
            return dark_regions, has_dark_edge
        
        # 存储所有区域的信息
        all_regions = []
        
        for i, contour in enumerate(contours):
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            
            if area < 100:  # 过滤掉太小的区域
                continue
            
            # 计算边界框
            x, y, w, h = cv2.boundingRect(contour)
            bbox = [x, y, x + w, y + h]
            
            # 计算轮廓中心
            M = cv2.moments(contour)
            if M["m00"] != 0:
                center_x = int(M["m10"] / M["m00"])
                center_y = int(M["m01"] / M["m00"])
            else:
                center_x = x + w // 2
                center_y = y + h // 2
            
            # 获取原始轮廓点
            polygon_points = []
            for point in contour:
                polygon_points.append([int(point[0][0]), int(point[0][1])])
            
            # 智能简化：根据轮廓长度和复杂度决定简化程度
            contour_length = cv2.arcLength(contour, True)
            if len(polygon_points) > 2000:  # 如果点数太多才进行简化
                # 使用更小的epsilon值来保持细致度
                epsilon = 0.0005 * contour_length  # 非常小的简化因子
                approx_polygon = cv2.approxPolyDP(contour, epsilon, True)
                polygon_points = []
                for point in approx_polygon:
                    polygon_points.append([int(point[0][0]), int(point[0][1])])
            elif len(polygon_points) > 500:  # 中等点数，轻微简化
                epsilon = 0.0002 * contour_length  # 极小的简化因子
                approx_polygon = cv2.approxPolyDP(contour, epsilon, True)
                polygon_points = []
                for point in approx_polygon:
                    polygon_points.append([int(point[0][0]), int(point[0][1])])
            # 如果点数少于500，保持原始轮廓点
            
            region_info = {
                'contour': contour,  # 原始轮廓
                'polygon': polygon_points,  # 简化后的多边形点
                'bbox': bbox,  # 边界框
                'area': int(area),  # 面积
                'center': [center_x, center_y],  # 中心点
                'confidence': 1.0,  # 分割模型通常输出确定性结果
                'region_id': i
            }
            
            all_regions.append(region_info)
        
        if not all_regions:
            return dark_regions, has_dark_edge
        
        # 按面积排序，找到最大的连通域
        all_regions.sort(key=lambda x: x['area'], reverse=True)
        largest_region = all_regions[0]
        
        # 检查最大区域是否超过面积阈值
        if largest_region['area'] >= adjusted_area_threshold:
            has_dark_edge = True
            # 只返回最大的区域
            dark_regions = [largest_region]
        else:
            has_dark_edge = False
            # 如果最大区域面积不够，返回空列表
            dark_regions = []
        
        return dark_regions, has_dark_edge
    
    def _save_result_image(self, original_image: Union[str, Image.Image, np.ndarray], 
                          result: Dict[str, Any], output_path: str):
        """
        保存带标注的结果图片
        
        Args:
            original_image: 原始图像
            result: 检测结果
            output_path: 输出路径
        """
        try:
            import matplotlib.pyplot as plt
            import matplotlib.patches as patches
            
            # 加载原始图像
            if isinstance(original_image, str):
                img = Image.open(original_image)
            elif isinstance(original_image, Image.Image):
                img = original_image
            elif isinstance(original_image, np.ndarray):
                if original_image.dtype != np.uint8:
                    original_image = (original_image * 255).astype(np.uint8)
                img = Image.fromarray(original_image)
            else:
                raise ValueError(f"不支持的图像类型: {type(original_image)}")
            
            # 创建图形
            fig, axes = plt.subplots(1, 2, figsize=(15, 6))
            
            # 显示原图
            axes[0].imshow(img)
            axes[0].set_title('原始图像')
            axes[0].axis('off')
            
            # 显示分割掩码
            seg_mask = result['segmentation_mask']
            axes[1].imshow(seg_mask, cmap='gray')
            axes[1].set_title(f'分割掩码 (检测到 {result["total_dark_regions"]} 个暗区)')
            axes[1].axis('off')
            
            # 在原图上绘制暗区边界框
            axes[0].imshow(img)
            for i, region in enumerate(result['dark_regions']):
                bbox = region['bbox']
                rect = patches.Rectangle(
                    (bbox[0], bbox[1]), 
                    bbox[2] - bbox[0], 
                    bbox[3] - bbox[1],
                    linewidth=2, 
                    edgecolor='red', 
                    facecolor='none'
                )
                axes[0].add_patch(rect)
                
                # 添加区域编号
                axes[0].text(bbox[0], bbox[1] - 5, f'暗区{i+1}', 
                           color='red', fontsize=10, fontweight='bold')
            
            plt.tight_layout()
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
            plt.close()
            
            print(f"结果图片已保存到: {output_path}")
            
        except Exception as e:
            print(f"保存结果图片失败: {str(e)}")
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        try:
            return self.infer.get_model_info()
        except Exception as e:
            return {'error': f"获取模型信息失败: {str(e)}"}
    
    def validate_model(self) -> bool:
        """
        验证模型是否有效
        
        Returns:
            bool: 模型是否有效
        """
        try:
            return self.infer.validate_model()
        except Exception as e:
            print(f"模型验证失败: {str(e)}")
            return False
    
    def get_task_type(self) -> str:
        """
        获取任务类型
        
        Returns:
            str: 任务类型
        """
        try:
            return self.infer.get_task_type()
        except Exception as e:
            return f"获取任务类型失败: {str(e)}"
